import os
os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'

from ultralytics import YOLO
import cv2
import matplotlib.pyplot as plt
from matplotlib import rcParams
import numpy as np

# 设置支持中文的字体
rcParams['font.family'] = 'sans-serif'
rcParams['font.sans-serif'] = ['SimHei']  # 使用中文字体

def adjust_brightness(image, factor):
    """调节图像亮度"""
    hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)  # 将图像转换为HSV模式
    h, s, v = cv2.split(hsv)
    v = np.clip(v * factor, 0, 255).astype(np.uint8)  # 调整亮度（v通道）
    adjusted_hsv = cv2.merge([h, s, v])
    adjusted_image = cv2.cvtColor(adjusted_hsv, cv2.COLOR_HSV2BGR)  # 转回BGR模式
    return adjusted_image

def infer(image, model, factor):
    """对输入图像在不同亮度下进行推理"""
    # 调节亮度
    bright_img = adjust_brightness(image, factor)

    # 转换为RGB模式（YOLO要求）
    img_rgb = cv2.cvtColor(bright_img, cv2.COLOR_BGR2RGB)

    # 执行推理
    results = model(img_rgb)

    # 可视化检测结果
    result_img = results[0].plot()  # 绘制检测结果
    plt.imshow(result_img)
    plt.title(f'推理结果 (亮度调整系数: {factor})')
    plt.axis('off')
    plt.show()

    # 输出检测的类别、置信度和边界框信息
    print(f"亮度调整系数: {factor}")
    for detection in results[0].boxes.data:
        x1, y1, x2, y2, score, class_id = detection.tolist()
        print(f"类别: {model.names[int(class_id)]}, 置信度: {score:.2f}, 边界框: ({x1}, {y1}, {x2}, {y2})")

def main():
    # 加载训练好的YOLO模型
    model = YOLO('./runs/detect/yolo_custom2/weights/best.pt')  # 使用训练后保存的模型

    # 加载图像
    image_path = './image.jpg'
    img = cv2.imread(image_path)

    # 对不同亮度条件的图像进行推理
    brightness_factors = [0.5, 0.75, 1.0, 1.25, 1.5]  # 定义亮度调整系数（<1为减弱，>1为增强）
    for factor in brightness_factors:
        infer(img, model, factor)

if __name__ == '__main__':
    main()
